from BigDataWeb.algorithm import Algorithm
import numpy as np
from keras.applications import inception_v3, resnet50, mobilenet
from keras.preprocessing.image import load_img
from keras.preprocessing.image import img_to_array
from keras.applications.imagenet_utils import decode_predictions
import os
import shutil
import BigDataWeb.config as config
from multiprocessing import Process, Manager


def runImageRecognitionModel(data):
    data_source_file_path = data["data_source_file_path"]
    
    inception_model = inception_v3.InceptionV3(weights='imagenet')
    resnet50_model = resnet50.ResNet50(weights='imagenet')
    mobilenet_model = mobilenet.MobileNet(weights='imagenet')
    # 以PIL格式加载图像，注意ImageNet模型只能处理224x224图片
    original = load_img(data_source_file_path, target_size=(224, 224))
    image_batch = np.expand_dims(img_to_array(original), axis=0)
    
    # ResNet50
    processed_image = resnet50.preprocess_input(image_batch.copy())
    predictions = resnet50_model.predict(processed_image)
    prediction_resnet = decode_predictions(predictions, top=1)
    label_resnet = prediction_resnet[0][0][1]
    possible_resnet = prediction_resnet[0][0][2]
    
    # MobileNet
    processed_image = mobilenet.preprocess_input(image_batch.copy())
    predictions = mobilenet_model.predict(processed_image)
    prediction_mobilnet = decode_predictions(predictions, top=1)
    label_mobilnet = prediction_mobilnet[0][0][1]
    possible_mobilnet = prediction_mobilnet[0][0][2]
    
    # InceptionV3
    # 初始网络的输入大小与其他网络不同。 它接受大小的输入（299,299）。
    original = load_img(data_source_file_path, target_size=(299, 299))
    image_batch = np.expand_dims(img_to_array(original), axis=0)
    processed_image = inception_v3.preprocess_input(image_batch.copy())
    predictions = inception_model.predict(processed_image)
    prediction_inception = decode_predictions(predictions, top=1)
    label_inception = prediction_inception[0][0][1]
    possible_inception = prediction_inception[0][0][2]
    
    data["label_resnet"] = label_resnet
    data["possible_resnet"] = possible_resnet
    data["label_mobilnet"] = label_mobilnet
    data["possible_mobilnet"] = possible_mobilnet
    data["label_inception"] = label_inception
    data["possible_inception"] = possible_inception


class ImageRecognition(Algorithm):
    label_resnet = ""
    possible_resnet = 0
    label_mobilnet = ""
    possible_mobilnet = 0
    label_inception = ""
    possible_inception = 0
    
    def readDataSouce(self, upload_file_path=""):
        # 读取数据源
        # 准备数据源文件
        if upload_file_path:
            # 手工上传
            self.data_source_name = "picture.jpg"
            self.data_source_file_path = os.path.join(self.work_folder_path, self.data_source_name)
            # 移动
            shutil.move(upload_file_path, self.data_source_file_path)
        else:
            # 本地模板
            self.data_source_file_path = os.path.join(self.work_folder_path, self.data_source_name)
            # 拷贝
            shutil.copyfile(os.path.join(config.sample_folder_path, self.data_source_name), self.data_source_file_path)

    def __init__(self):
        Algorithm.__init__(self)
        self.algorithm_name = "深度学习-图像识别"
        self.ipynb_template_name = "image_recognition-template.ipynb"

    def implent(self):
        Algorithm.implent(self)
        
        with Manager() as manager:
            data = manager.dict()
            data["data_source_file_path"] = self.data_source_file_path
            process = Process(target=runImageRecognitionModel, args=(data,))
            process.start()
            process.join()
            
            self.label_resnet = data["label_resnet"]
            self.possible_resnet = data["possible_resnet"]
            self.label_mobilnet = data["label_mobilnet"]
            self.possible_mobilnet = data["possible_mobilnet"]
            self.label_inception = data["label_inception"]
            self.possible_inception = data["possible_inception"]
